import os
import sys
import webview
import cv2
import numpy as np
import base64
import io
import time
import psutil
import json
from PIL import Image

# 添加项目根目录到系统路径
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.append(project_root)

# 导入算法模块
from algorithms.tradition import stitch_images as stitch_images_tradition
from algorithms.udis2 import stitch_images_udis2, cleanup_temp_files # UDIS2 import
# from algorithms.udis2_refactored.stitcher import stitch_images as stitch_images_udis2 # Refactored UDIS2 import

class ImageStitchingAPI:
    def __init__(self):
        self.img1 = None
        self.img2 = None
        self.result_image = None
        self.algorithm = "tradition"  # 默认使用传统算法
        self.processing_steps = []  # 用于存储UDIS2算法的中间步骤
        self.performance_data = {}  # 用于存储性能数据

    def upload_image(self, image_data, index):
        """接收前端上传的图片数据"""
        try:
            print(f"开始处理上传的图片 {index}，数据长度: {len(image_data)}")
            print(f"图片数据前50个字符: {image_data[:50]}...")
            
            # 解码base64图片数据
            image_data = image_data.split(',')[1]
            image_bytes = base64.b64decode(image_data)
            
            print(f"图片 {index} 解码后的字节长度: {len(image_bytes)}")
            
            # 转换为OpenCV格式
            nparr = np.frombuffer(image_bytes, np.uint8)
            img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            
            print(f"图片 {index} 转换为OpenCV格式后的形状: {img.shape if img is not None else 'None'}")
            
            if index == 1:
                self.img1 = img
                print("已保存为图片1")
            else:
                self.img2 = img
                print("已保存为图片2")
                
            return {"success": True, "message": f"图片 {index} 上传成功"}
        except Exception as e:
            print(f"图片 {index} 上传失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return {"success": False, "message": f"图片上传失败: {str(e)}"}

    def set_algorithm(self, algorithm):
        """设置拼接算法"""
        self.algorithm = algorithm
        return {"success": True, "message": f"已选择 {algorithm} 算法"}

    def stitch(self):
        """执行图像拼接"""
        print("开始执行图像拼接...")
        if self.img1 is None or self.img2 is None:
            print("拼接失败: 图片未完全上传")
            print(f"图片1状态: {'已上传' if self.img1 is not None else '未上传'}")
            print(f"图片2状态: {'已上传' if self.img2 is not None else '未上传'}")
            return {"success": False, "message": "请先上传两张图片"}
        
        try:
            # 记录开始时间和初始资源使用情况
            start_time = time.time()
            initial_memory = psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024  # MB
            initial_cpu = psutil.cpu_percent(interval=0.1)
            
            # 根据选择的算法执行拼接
            if self.algorithm == "tradition":
                print("执行传统算法")
                print(f"图片1形状: {self.img1.shape}")
                print(f"图片2形状: {self.img2.shape}")
                self.result_image = stitch_images_tradition(self.img1, self.img2)
                
                # 记录结束时间和最终资源使用情况
                end_time = time.time()
                final_memory = psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024  # MB
                final_cpu = psutil.cpu_percent(interval=0.1)
                
                # 存储性能数据
                self.performance_data = {
                    "execution_time": end_time - start_time,
                    "memory_usage": final_memory - initial_memory,
                    "cpu_usage": (initial_cpu + final_cpu) / 2
                }
                
                if self.result_image is None:
                    print("拼接失败，可能是匹配点不足")
                    return {"success": False, "message": "拼接失败，可能是匹配点不足"}
                
                print(f"拼接成功，结果图像形状: {self.result_image.shape}")
                
                # 保存结果图像
                result_path = os.path.join(project_root, "static", "results", "result.jpg")
                os.makedirs(os.path.dirname(result_path), exist_ok=True)
                cv2.imwrite(result_path, self.result_image)
                print(f"结果图像已保存到{result_path}")
                
                # 将结果图像转换为base64
                _, buffer = cv2.imencode('.jpg', self.result_image)
                img_base64 = base64.b64encode(buffer).decode('utf-8')
                
                return {
                    "success": True, 
                    "message": "拼接成功", 
                    "image": f"data:image/jpeg;base64,{img_base64}",
                    "performance": self.performance_data
                }
            elif self.algorithm == "udis2":
                print("执行UDIS2算法")
                print(f"图片1形状: {self.img1.shape}")
                print(f"图片2形状: {self.img2.shape}")
                
                # 调用UDIS2算法
                self.result_image, self.processing_steps, self.performance_data = stitch_images_udis2(self.img1, self.img2)
                
                if self.result_image is None:
                    print("UDIS2拼接失败")
                    return {"success": False, "message": "UDIS2拼接失败，请检查图像是否合适"}
                
                print(f"UDIS2拼接成功，结果图像形状: {self.result_image.shape}")
                
                # 保存结果图像
                result_path = os.path.join(project_root, "static", "results", "udis2_result.jpg")
                os.makedirs(os.path.dirname(result_path), exist_ok=True)
                cv2.imwrite(result_path, self.result_image)
                
                # 将结果图像转换为base64
                _, buffer = cv2.imencode('.jpg', self.result_image)
                img_base64 = base64.b64encode(buffer).decode('utf-8')
                
                # 处理中间步骤图像
                steps_data = []
                for step in self.processing_steps:
                    if os.path.exists(step["path"]):
                        with open(step["path"], "rb") as img_file:
                            img_data = base64.b64encode(img_file.read()).decode('utf-8')
                            steps_data.append({
                                "name": step["name"],
                                "image": f"data:image/jpeg;base64,{img_data}"
                            })
                
                return {
                    "success": True, 
                    "message": "UDIS2拼接成功", 
                    "image": f"data:image/jpeg;base64,{img_base64}",
                    "performance": self.performance_data,
                    "steps": steps_data
                }
            else:
                return {"success": False, "message": f"未知算法: {self.algorithm}"}
        except Exception as e:
            print(f"拼接过程中出错: {str(e)}")
            import traceback
            traceback.print_exc()
            return {"success": False, "message": f"拼接过程中出错: {str(e)}"}

    def get_performance_data(self):
        """获取性能数据"""
        return self.performance_data

    def reset(self):
        """重置状态"""
        self.img1 = None
        self.img2 = None
        self.result_image = None
        self.processing_steps = []
        self.performance_data = {}
        try:
            cleanup_temp_files() # 调用清理函数
            print("已清理UDIS2临时文件")
        except Exception as e:
            print(f"清理临时文件时出错: {str(e)}")
        return {"success": True, "message": "已重置"}

    def save_result(self):
        """打开保存对话框让用户选择路径并保存结果图像"""
        if self.result_image is None:
            return {"success": False, "message": "没有可保存的结果图像"}
        
        try:
            # 打开文件保存对话框
            file_path = webview.windows[0].create_file_dialog(
                webview.SAVE_DIALOG, 
                directory=os.path.expanduser('~'), # 使用用户主目录作为默认目录
                save_filename='panorama.jpg',
                file_types=('JPEG 文件 (*.jpg;*.jpeg)', 'PNG 文件 (*.png)')
            )
            
            # 用户取消了对话框
            if not file_path:
                print("用户取消了保存")
                return {"success": False, "message": "用户取消了保存"}
            
            # 如果是元组（Windows平台返回元组），取第一个元素
            if isinstance(file_path, tuple):
                file_path = file_path[0]
            
            print(f"准备保存图像到: {file_path}")
            # 获取文件扩展名
            _, ext = os.path.splitext(file_path)
            
            # 编码图像
            # OpenCV的imencode函数需要一个以点开头的扩展名，例如 '.jpg'
            encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 90] # 默认JPEG质量
            if ext.lower() == '.png':
                encode_param = [int(cv2.IMWRITE_PNG_COMPRESSION), 3] # 默认PNG压缩级别

            success, buffer = cv2.imencode(ext, self.result_image, encode_param)
            
            if success:
                try:
                    with open(file_path, 'wb') as f:
                        f.write(buffer)
                    print(f"图像已成功保存到 {file_path}")
                    return {"success": True, "message": f"图像已保存到 {file_path}"}
                except Exception as e_write:
                    print(f"写入文件时出错: {file_path}, 错误: {str(e_write)}")
                    return {"success": False, "message": f"保存图像失败: {str(e_write)}"}
            else:
                print(f"图像编码失败: {file_path}")
                return {"success": False, "message": f"保存图像失败: 图像编码失败"}
        except Exception as e:
            print(f"保存过程中出错: {str(e)}")
            import traceback
            traceback.print_exc()
            return {"success": False, "message": f"保存失败: {str(e)}"}


def main():
    # 确保静态资源目录存在
    static_dir = os.path.join(project_root, "static")
    css_dir = os.path.join(static_dir, "css")
    js_dir = os.path.join(static_dir, "js")
    results_dir = os.path.join(static_dir, "results")
    
    os.makedirs(css_dir, exist_ok=True)
    os.makedirs(js_dir, exist_ok=True)
    os.makedirs(results_dir, exist_ok=True)
    
    # HTML文件路径
    html_path = os.path.join(static_dir, "index.html")
    
    # 创建API实例
    api = ImageStitchingAPI()
    
    # 创建窗口
    window = webview.create_window(
        '图像拼接工具', 
        html_path, 
        js_api=api,
        width=1000, 
        height=800, 
        resizable=True,
        text_select=True
    )
    
    # 启动应用
    webview.start(debug=False)


if __name__ == '__main__':
    main()
